package com.cowmooed.service.impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import com.cowmooed.service.PromptService;
import org.springframework.stereotype.Service;

import java.util.Arrays;

@Service
public class PromptServiceImpl implements PromptService {

    public static GenerationResult callWithMessage(String prompt) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你的任务是对用户的提示词进行优化，要求：\" +\n" +
                        "\"1. 主体描述 - 主体描述是对主体对象的描述，建议以单主体为主\" +\n" +
                        "\"2. 特征描述 - 可同时添加对主体形状、颜色、材质、风格等细节描述，如大眼睛、蓝色等，不需要添加常用生图的 prompt 技巧，例如 negative prompt，4k 等词\" +\n" +
                        " \"3. 推荐格式 - {量词，建议为1}+{主体}，{可选：详细描述}，{可选：风格}\"" +
                         " \"举例: [\"一只深棕色、大眼睛的猫头鹰\", \"一件黄色的连衣裙，上面印有黄色和绿色的图案，呈现真实摄影风格\"" +
                         "提示词可能会存在一些用户随意输入的内容，需要你提示用户，并且给用户一些建议，" +
                        "最后返回 { \"success\": \"\", \"prompt\": \"\", \"suggestion\": \"\" }格式 success代表用户输入是否合法，取值true or false || prompt在用户提示词合法时返回, 否则返回空字符串 || suggestion 在用户提示词不合法时返回，否则返回空字符串"
                                )

                .build();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content("[提示词: " + prompt + "] 请你对提示词进行优化并且返回新的提示词")
                .build();
        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                // 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model("qwen-plus")
                .messages(Arrays.asList(systemMsg, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        return gen.call(param);
    }

    @Override
    public String OptimizePrompt(String prompt) {
        try {
            GenerationResult result = callWithMessage(prompt);
            System.out.println(JsonUtils.toJson(result));
            return result.getOutput().getChoices().get(0).getMessage().getContent();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            // 使用日志框架记录异常信息
            System.err.println("An error occurred while calling the generation service: " + e.getMessage());
        }
        return "{ \"success\": true, \"prompt\": \"" + "\"" + prompt  +"\"}";
    }
}
